from kea2.hypiumDriver import HDriver
import json
import os
import re
from pathlib import Path


def get_hierarchy():
    d = HDriver.getScriptDriver()
    r = d.hdc("shell uitest dumpLayout")
    remote_dir = r.split(":")[-1]
    file_name = Path(remote_dir).name
    d.pull_file(remote_dir, ".")
    with open(file_name, "r") as fp:
        hierarchy = json.load(fp)
    os.remove(file_name)
    return hierarchy


def traverse_uiTree(node, hook, skip_system=True):
    """Recursively traverse the UI tree and apply the hook function to each node.
    Args:
        node (Dict): The current node to check.
        hook (Callable): A function that takes node attributes and returns True if it matches.
        skip_system (bool): If True, skip nodes with a specific bundleName.
    Returns:
        bool: True if the hook function returns True for any node, False otherwise.
    """
    if skip_system and node["attributes"].get("bundleName") in {
        "com.ohos.sceneboard"
    }:
        node["children"] = []
        return False
    if hook(node["attributes"]):
        print(node["attributes"])
        return True
    for child in node["children"]:
        if traverse_uiTree(child, hook):
            return True
    return False


def get_node(node, hook):
    """Recursively search for a node that matches the hook condition.
    Args:
        node (Dict): The current node to check.
        hook (Callable): A function that takes node attributes and returns True if it matches.
    """
    if hook(node["attributes"]):
        return node
    for child in node["children"]:
        if (n := get_node(child, hook)):
            return n


def get_bounds(node_attrs):
    """ Extract the bounds of a node from its attributes.
    Args:
        node_attrs (Dict): The attributes of the node.
    Returns:
        List[int]: A list containing the bounds in the format [left, top, right, bottom].
    e.g. [0, 0, 1080, 1920]
    """
    if node_attrs.get("attributes", None) is not None:
        node_attrs = node_attrs["attributes"]
    
    pattern = re.compile(r"\[(-?\d+),(-?\d+)\]\[(-?\d+),(-?\d+)\]")
    m = re.match(pattern, node_attrs["bounds"])
    bounds = [int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4))]

    return bounds


class TreeInfo:
    """Class to hold information about the UI tree.
    """
    _width = None
    _height = None
    
    def __init__(self, root):
        self.root = root
    
    @property
    def width(self):
        if self._width is None:
            self._cal_tree_info()
        return self._width
    
    @property
    def height(self):
        if self._height is None:
            self._cal_tree_info()
        return self._height
    
    def _cal_tree_info(self):
        self._height = height = 0
        self._width = width = 0
        
        def traverse(node):
            nonlocal height, width
            height += 1
            self._height = max(self._height, height)
            self._width = max(self._width, len(node["children"]))
            for child in node["children"]:
                traverse(child)
        
        traverse(self.root)


if __name__ == "__main__":
    # Example usage
    HDriver.initializer("23E0224307001843")
    driver = HDriver.getInstance()
    hierarchy = get_hierarchy()
    tree_info = TreeInfo(hierarchy)
    
    print(f"Tree Width: {tree_info.width}, Height: {tree_info.height}")